import csv
from math import nan
import psycopg2

g_csv_file_path = 'D:\\wuxi data\\PMO\\三资\\报告\\数据处理\\工时导入.csv'
# 假设从 config 复制的数据库参数初始化逻辑
def init_db_params():
    return {
        # 这里填写具体的数据库参数，示例如下
        'host': 'localhost',
        'port': '5432',
        'database': 'sanzi',
        'user': 'postgres',
        'password': 'Staff1234'
    }

def handle_t_work_hour_records(conn, cursor):
    # 检查表是否存在
    v_check_table_query = "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 't_work_hour_records')"
    cursor.execute(v_check_table_query)
    v_table_exists = cursor.fetchone()[0]

    if not v_table_exists:
        # 创建表
        v_create_table_query = '''
        CREATE  TABLE  t_work_hour_records (
             record_id SERIAL PRIMARY  KEY,                 -- 自增主键
             organization TEXT,                            -- 所属组织
             organization_code VARCHAR (50),                -- 所属组织编号
             employee_name TEXT,                           -- 姓名
             employee_id VARCHAR (50),                      -- 人员编号
             department TEXT,                              -- 部门
             task_name TEXT,                               -- 任务名称
             project_manager TEXT,                         -- 归属项目经理
             project_name TEXT,                            -- 归属项目
             project_code VARCHAR (50),                     -- 归属项目编号
             task_type TEXT,                               -- 任务类型
             time_report_time TIMESTAMP,                   -- 工时填报时间
             time_belong_date DATE,                        -- 工时归属日期
             reported_hours NUMERIC (6,2),                  -- 对应填报工时（小时）
             approved_hours NUMERIC (6,2),                  -- 实际认定工时（小时）
             approval_time TIMESTAMP,                      -- 认定时间
             other_notes TEXT,                             -- 其他说明
             processing_notes TEXT,                        -- 处理说明
             from_source TEXT,                             -- 来源
             created_at TIMESTAMP  DEFAULT  CURRENT_TIMESTAMP -- 记录创建时间
         );
         -- 添加常用查询字段的索引
         CREATE  INDEX  idx_employee_id ON  t_work_hour_records (employee_id);
         CREATE  INDEX  idx_project_code ON  t_work_hour_records (project_code);
         CREATE  INDEX  idx_time_belong_date ON  t_work_hour_records (time_belong_date);
        '''
        cursor.execute(v_create_table_query)
        print('表 t_work_hour_records 创建成功')

    if v_table_exists:
        # 清理表中全部数据
        v_delete_query = "TRUNCATE TABLE t_work_hour_records RESTART IDENTITY"
        cursor.execute(v_delete_query)
        conn.commit()
        print('表 t_work_hour_records 数据清理成功')
        
# 从 db_setup 复制的工时表建表初始化逻辑
def setup_database(conn, cursor):
    try:
        handle_t_work_hour_records(conn, cursor)

    except Exception as e:
        print(f'发生错误: {e}')
        raise e  # 重新抛出异常，以便
    


# 封装工时导入方法
def import_to_t_work_hour_records_from_csv(conn, cursor):
    # 打开CSV文件并导入数据
    with open(g_csv_file_path, 'r', newline='', encoding='utf-8') as file:
        v_reader = csv.reader(file, delimiter=',')
        next(v_reader)  # 跳过标题行
        v_rows = [row for row in v_reader]
        # 对 reported_hours 字段做四舍五入处理，保留两位小数
        v_time_report_time_index = 10  # 根据插入语句中字段顺序确定 time_report_time 索引
        v_time_belong_date_index = 11  # 根据插入语句中字段顺序确定 time_belong_date 索引
        v_reported_hours_index = 12  # 根据插入语句中字段顺序确定 reported_hours 索引
        v_approved_hours_index = 13  # 根据插入语句中字段顺序确定 approved_hours 索引
        v_approval_time_index = 14  # 根据插入语句中字段顺序确定 approval_time 索引
        v_batch_size = 1000
        for v_i in range(0, len(v_rows), v_batch_size):
            v_batch = v_rows[v_i:v_i + v_batch_size]
            for v_row in v_batch:
                if not v_row[v_time_report_time_index]:
                    v_row[v_time_report_time_index] = None
                if not v_row[v_time_belong_date_index]:
                    v_row[v_time_belong_date_index] = None
                if v_row[v_reported_hours_index]:
                    v_row[v_reported_hours_index] = round(float(v_row[v_reported_hours_index]), 2)
                else:
                    v_row[v_reported_hours_index] = None
                if v_row[v_approved_hours_index]:
                    v_row[v_approved_hours_index] = round(float(v_row[v_approved_hours_index]), 2)
                else:
                    v_row[v_approved_hours_index] = None
                if not v_row[v_approval_time_index]:
                    v_row[v_approval_time_index] = None
            # print('即将插入数据库的CSV批次内容：', batch)
            v_insert_query = '''
            INSERT INTO t_work_hour_records (
                organization, organization_code, employee_name, employee_id, department,
                task_name, project_manager, project_name, project_code, task_type,
                time_report_time, time_belong_date, reported_hours, approved_hours,
                approval_time, other_notes, processing_notes, from_source
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            '''
            try:
                cursor.executemany(v_insert_query, v_batch)
                conn.commit()
                v_end_index = v_i + min(len(v_batch), v_batch_size) - 1
                print('批次数据导入成功，批次范围：', v_i + 1, ' - ', v_end_index + 1)
            except Exception as e:
                conn.rollback()
                v_error_row_index = 0
                for v_index, v_row in enumerate(v_batch):
                    try:
                        cursor.execute(v_insert_query, v_row)
                    except Exception as inner_e:
                        v_error_row_index = v_i + v_index
                        v_end_index = v_i + min(len(v_batch), v_batch_size) - 1
                        print(f'批次数据导入失败，当前批次: {v_i // v_batch_size + 1}，开始记录数: {v_i + 1}，结束记录数: {v_end_index + 1}，错误数据所在行: {v_error_row_index + 2}（包含标题行），错误信息: {inner_e}')
                        print('错误的行数据如下：', v_row)
                        break
                conn.rollback()

# 主方法
def main():
    try:
        db_params = init_db_params()
        g_conn = psycopg2.connect(**db_params)
        g_cursor = g_conn.cursor()
        # 调用数据库设置函数
        setup_database(g_conn, g_cursor)

        # 调用封装的导入方法
        import_to_t_work_hour_records_from_csv(g_conn, g_cursor)

    except Exception as e:
        print(f'发生错误，错误信息: {e}')
        if g_conn:
            g_conn.rollback()
    finally:
        if g_conn:
            g_cursor.close()
            g_conn.close()

if __name__ == '__main__':
    main()